import cPickle as pickle
import matplotlib
#matplotlib.use("TkAgg")
from pylab import *
from numpy import array, vstack, sqrt
import os
from utils import is_correct, confidence_interval, meas_fid, onlyfiles

f_real = lambda cpx_num: cpx_num.real 

def plot_fid_v_t(dir_name):
    """
    For all .pkl files in a directory, unpickles, looks for tmeas_off
    and s, overlap_p, overlap_m lists, then runs meas_fid on the list 
    and plots the resulting pair of lists t_list and fid_list.
    """
    
    t_list = []
    fid_list = []
    fil_fid_list = []
    ci_list = []
    fil_ci_list = []
    upr_list = []
    fil_upr_list = []
    lwr_list = []
    fil_lwr_list = []
    
    for phylo in onlyfiles(dir_name):
        with open(dir_name + '/' + phylo) as phil:
            data_dict = pickle.load(phil)
        t_list.append(data_dict['tmeas_off'])    
        fid_stats = meas_fid(map(f_real, data_dict['s']), data_dict['overlap_p'], data_dict['overlap_m'])
        fil_fid_stats = meas_fid(map(f_real, data_dict['s_filtered']), data_dict['overlap_p'], data_dict['overlap_m'])
        fid_list.append(fid_stats[0])            
        fil_fid_list.append(fil_fid_stats[0])    
        ci_list.append(list(fid_stats[1]))
        lwr_list.append(fid_stats[1][0])
        upr_list.append(fid_stats[1][1])
        fil_ci_list.append(list(fil_fid_stats[1]))
        fil_upr_list.append(fil_fid_stats[1][1])
        fil_lwr_list.append(fil_fid_stats[1][0])
    
    yerr_arr = vstack([array(upr_list) - array(fid_list), array(fid_list) - array(lwr_list)])
    fil_yerr_arr = vstack([array(fil_upr_list) - array(fil_fid_list), array(fil_fid_list) - array(fil_lwr_list)])
    
    #errorbar(t_list, fid_list, fmt='.', yerr = array(ci_list))
    #errorbar(t_list, fil_fid_list, fmt='.', yerr = array(fil_ci_list))
    errorbar(t_list, fid_list, fmt='o', yerr = array(upr_list) - array(fid_list))
    errorbar(t_list, fil_fid_list, fmt='o', yerr = array(fil_upr_list) - array(fil_fid_list))
    show()

def _plot_files(file_list, indy_par='total_height', extra_sqrt=False):
    x_list = []
    y_list = []
    yerr_list = [[],[]]
    for phylo in file_list:
        if phylo != "README":
            with open(phylo) as phil:
                data_dict = pickle.load(phil)
            x_list.append(data_dict[indy_par])
            if extra_sqrt:
                y=sqrt(data_dict['fid'])
            else:
                y=data_dict['fid']
            y_list.append(y)
            confidence_interval = data_dict['confidence']
            yerr_list[0].append(abs(confidence_interval[0] - y))
            yerr_list[1].append(abs(confidence_interval[1] - y))
    errorbar(x_list, y_list, yerr_list, fmt='o')
    
def optimal_control_fid_plot(dir_name='saved_traj',
                            indy_par='total_height', 
                            extra_sqrt=False):
    """
    Searches the given directory for files.
    They'd better contain pickled dicts. It unpickles them, looks for 
    an independent value to put on the x axis using the input key 
    `indy_par`, then produces an error bar plot with the point in black
    if the number of iterations is below `iter_max`, and in red if it's
    above or equal. 
    """
    _plot_files(onlyfiles(dir_name), indy_par, extra_sqrt)
    show()

def two_param_error_plot(dir_name, x_param = 'total_height',
                        series_param = 'y_offset', extra_sqrt=False):
    """
    Looks for filenames that match a given format, then puts them 
    through _plot_files (this may become a pattern later). These files
    are to be named with two floating-point parameters separated by 
    underscores. These parameters are supposed to be at the end, so 
    for example, the file a_b_c_0.1_1.0 would be parsed as having
    an x_param of 0.1 and a series_param of 1.0. 
    """
    #First, collect a list of series_param values:
    series_param_vals = []
    big_file_list = onlyfiles(dir_name)
    for filename in big_file_list:
        potential_val = filename.split('_')[-1]
        if not(potential_val in series_param_vals):
            series_param_vals.append(potential_val)
    
    for series_param_val in series_param_vals:
        series_param_pred = lambda f: f.split('_')[-1] == series_param_val
        _plot_files(filter(series_param_pred, big_file_list))

    show()


def plot_fid_v_t_and_eps(dir_name, filtered=True, xtra_nms=None, xtra_vals=None):
    """
    For all .pkl files in a directory, unpickles, looks for tmeas_off, 
    max_eps, s, s_filtered, overlap_p, overlap_m lists, then runs 
    meas_fid on the list and plots points on the tmeas_off, max_eps 
    grid, annotated with the results of meas_fid. We allow the entry of
    extra parameters, whose names and values are given by xtra_nms and 
    xtra_vals. 
    """
    
    t_list = []
    fid_list = []
    fil_fid_list = []

    for phylo in onlyfiles(dir_name):
        with open(dir_name + '/' + phylo) as phil:
            data_dict = pickle.load(phil)

        #print data_dict.keys()            

        if filtered:
            fid, low_up = meas_fid(map(f_real, data_dict['s_filtered']), data_dict['overlap_p'], data_dict['overlap_m'])        
        else:
            fid, low_up = meas_fid(map(f_real, data_dict['s']), data_dict['overlap_p'], data_dict['overlap_m'])        
        if xtra_nms is None:
            scatter(data_dict['tmeas_off'], data_dict['max_eps'], marker='o')
            annotate(r'${:.3}'.format(fid)+'_{'+'{:.3}'.format(low_up[0])+r'}^{'+'{:.3}'.format(low_up[1])+r'}$',
                    xy = (data_dict['tmeas_off'], data_dict['max_eps']), 
                    textcoords='offset points')
        elif all([abs(data_dict[name] - val)<10**-12 for name, val in zip(xtra_nms, xtra_vals)]):
            scatter(data_dict['tmeas_off'], data_dict['max_eps'], marker='o')
            annotate(r'${:.3}'.format(fid)+'_{'+'{:.3}'.format(low_up[0])+r'}^{'+'{:.3}'.format(low_up[1])+r'}$',
                    xy = (data_dict['tmeas_off'], data_dict['max_eps']), 
                    textcoords='offset points')
            
    show()


def plot_overlap_vs_s(filename, filtered='t', pms = 'p'):
    """
    Opens an output file, unpickles, plots either the filtered or 
    unfiltered measurement signal, or both, versus overlap_p,
    overlap_m, or the sum of the two, depending on whether filtered = 
    't', 'f', or 'b', and pms = 'p', 'm' or 's', respectively.
    """
    with open(filename, 'r') as phil:
        data_dct = pickle.load(phil)
    
    titelist = ['Overlap with','vs.', 'measurement signal']
    #Decide x variable, x label and part of title based on filtered:    
    if filtered == 't':
        titelist[2] = 'filtered measurement signal'
        indep = data_dct['s_filtered']
        x_label = "filtered measurement signal (dimensionless)"
        stl_str = "b."
    elif filtered == 'f':
        indep = data_dct['s']
        x_label = "unfiltered measurement signal"
        stl_str = "k."
    else:
        raise ValueError("Input variable 'filtered' must be Boolean;"+\
                            "you entered {0}".format(filtered))
    
    if pms == 'p':
        dep = data_dct['overlap_p']
        titelist.insert(1, 'plus state')
        y_label = ('overlap with plus state')
    elif pms == 'm':
        dep = data_dct['overlap_m']
        titelist.insert(1, 'minus state')
        y_label = ('overlap with minus state')
    elif pms == 's':
        dep = array(data_dct['overlap_p']) + \
                array(data_dct['overlap_m'])
        titelist.insert(1, 'plus/minus space')
        y_label = ('overlap with plus/minus subspace')
    else:
        raise ValueError("pms must be one of ['p', 'm', 's'], you"+\
                            " entered {0}".format(pms))
    
    #plot(indep, dep, stl_str)
    plot(data_dct['s_filtered'], dep, 'b.')
    plot(data_dct['s'], dep, 'k.')    
    xlabel(x_label)
    ylabel(y_label)
    #xlim([-80,80])
    #ylim([0.2,0.8])
    title(' '.join(titelist))    
    show()

def pause_plots(plot_cmd, **kwargs):
    for filename in onlyfiles(os.getcwd()):
        clf()        
        plot_cmd(filename, **kwargs)
        pause(1)
        
        


















